#include "Fit/RooFitBuilder.h"
#include "File/BOSSOutputLoader.h"
#include "Functions/Error.h"
#include <list>
#include <string>
using namespace RooFit;
using namespace BOSS_Afterburner::Error;
using namespace std;

Int_t RooFitBuilder::fID = 0;

RooRealVar* RooFitBuilder::BuildVar(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  if(!input["value"].IsScalar()) EXCEPTION("YAML variable definition does not contain a value");
  auto   value = input["value"].as<Double_t>();
  string name  = ExtractName(input);
  string unit;
  if(!!input["unit"] && input["unit"].IsScalar()) unit = input["unit"].as<string>();
  const auto& range = input["range"];
  if(!!range && range.size() == 2)
    return new RooRealVar(name.c_str(), name.c_str(), value, range[0].as<Double_t>(),
                          range[1].as<Double_t>(), unit.c_str());
  return new RooRealVar(name.c_str(), name.c_str(), value, unit.c_str());
}

RooArgusBG* RooFitBuilder::BuildArgus(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name = ExtractName(input);
  auto m0   = BuildVar(input["m0"]);
  auto c    = BuildVar(input["c"]);
  auto p    = BuildVar(input["p"]);
  return new RooArgusBG(name.c_str(), name.c_str(), fVar, *m0, *c, *p);
}

RooArgusBESIII* RooFitBuilder::BuildArgusBESIII(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name = ExtractName(input);
  auto m0   = BuildVar(input["m0"]);
  auto c    = BuildVar(input["c"]);
  auto p    = BuildVar(input["p"]);
  return new RooArgusBESIII(name.c_str(), name.c_str(), fVar, *m0, *c, *p);
}

RooBreitWigner* RooFitBuilder::BuildBreitWigner(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name  = ExtractName(input);
  auto mean  = BuildVar(input["mean"]);
  auto sigma = BuildVar(input["sigma"]);
  return new RooBreitWigner(name.c_str(), name.c_str(), fVar, *mean, *sigma);
}

RooAddPdf* RooFitBuilder::BuildGaussians(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  if(!input["sigmas"].IsSequence()) EXCEPTION("Does not contain sigma list");

  auto mean = BuildVar(input["mean"]);
  if(!mean) EXCEPTION("Mean has not been defined");

  auto        gaussians = new RooArgList;
  auto        coeff     = new RooArgList;
  const auto& sigmadefs = input["sigmas"];
  const auto& ratiodefs = input["ratios"];

  Int_t n = sigmadefs.size();
  Int_t r = 0;
  if(ratiodefs) r = input["ratios"].size();
  for(Int_t i = 0; i < n; ++i)
  {
    auto sigma = BuildVar(sigmadefs[i]);
    sigma->setUnit(mean->getUnit());
    sigma->SetName(Form("#sigma_{gauss%d}", i + 1));
    const char* rationame = Form("#it{N}_{gauss%d}", i + 1);
    RooRealVar* ratio     = nullptr;
    if(i < r)
      ratio = BuildVar(ratiodefs[i]);
    else
      ratio = new RooRealVar(rationame, rationame, .1, 0, 1);
    ratio->SetName(rationame);
    string name = Form("Gaussian_{%d}", i);
    gaussians->add(*(new RooGaussian(name.c_str(), name.c_str(), fVar, *mean, *sigma)));
    if(i < n - 1) coeff->add(*ratio);
  }
  return AddPdfs(gaussians, coeff);
}

RooGaussian* RooFitBuilder::BuildGaussian(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name  = ExtractName(input);
  auto mean  = BuildVar(input["mean"]);
  auto sigma = BuildVar(input["sigma"]);
  if(!mean || !sigma) EXCEPTION("Nullptrs in variables");
  return new RooGaussian(name.c_str(), name.c_str(), fVar, *mean, *sigma);
}

RooBreitWignerWave* RooFitBuilder::BuildPwave(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name  = ExtractName(input);
  auto mean  = BuildVar(input["mean"]);
  auto sigma = BuildVar(input["sigma"]);
  auto l     = new RooRealVar("l", "l", 1.);
  return new RooBreitWignerWave(name.c_str(), name.c_str(), fVar, *mean, *sigma, *l);
}

RooFFTConvPdf* RooFitBuilder::BuildConvolutedPwave(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  const auto& pwaveDef = input["P-wave"];
  const auto& gaussDef = input["Gaussian"];
  if(!pwaveDef || !pwaveDef.IsMap()) EXCEPTION("Missing key \"P-wave\"");
  if(!gaussDef || !gaussDef.IsMap()) EXCEPTION("Missing key \"Gaussian\"");
  auto pwave = BuildPwave(pwaveDef);
  auto gauss = BuildGaussian(gaussDef);
  return ConvolutePdf(pwave, gauss);
}

RooPolynomial* RooFitBuilder::BuildPolynomial(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  const auto& labelDef = input["label"];
  const auto& paramDef = input["parameters"];
  if(!paramDef || !paramDef.IsSequence() || !paramDef.size())
    EXCEPTION("Invalid sequence \"parameters\"");

  string label = "pol";
  if(!!labelDef && labelDef.IsScalar()) label = labelDef.as<string>();
  auto pars = new RooArgList;
  for(Int_t i = 0; i < paramDef.size(); ++i)
  {
    auto p = BuildVar(paramDef[i]);
    if(!p) continue;
    p->SetName(Form("#it{p}_{%d}", i + 1));
    pars->add(*p);
  }
  return new RooPolynomial(label.c_str(), label.c_str(), fVar, *pars);
}

RooChebychev* RooFitBuilder::BuildChebychev(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  const auto& labelDef = input["label"];
  const auto& paramDef = input["parameters"];
  if(!paramDef || !paramDef.IsSequence() || !paramDef.size())
    EXCEPTION("Invalid sequence \"parameters\"");

  string label = "cheb";
  if(!!labelDef && labelDef.IsScalar()) label = labelDef.as<string>();
  auto pars = new RooArgList;
  for(Int_t i = 0; i < paramDef.size(); ++i)
  {
    auto p = BuildVar(paramDef[i]);
    if(!p) continue;
    p->SetName(Form("#it{p}_{%d}", i + 1));
    pars->add(*p);
  }
  return new RooChebychev(label.c_str(), label.c_str(), fVar, *pars);
}

RooTruncatedPolynomial* RooFitBuilder::BuildTruncatedPolynomial(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name = ExtractName(input);
  auto m0   = BuildVar(input["m0"]);
  auto a    = BuildVar(input["a"]);
  auto b    = BuildVar(input["b"]);
  auto c    = BuildVar(input["c"]);
  return new RooTruncatedPolynomial(name.c_str(), name.c_str(), fVar, *m0, *a, *b, *c);
}

RooVoigtian* RooFitBuilder::BuildVoigtian(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  auto name  = ExtractName(input);
  auto mean  = BuildVar(input["mean"]);
  auto width = BuildVar(input["width"]);
  auto sigma = BuildVar(input["sigma"]);
  return new RooVoigtian(name.c_str(), name.c_str(), fVar, *mean, *width, *sigma);
}

RooHistPdf* RooFitBuilder::BuildHistPdf(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  const auto& fileDef = input["Input file or directory"];
  const auto& plotDef = input["Plot definition"];
  if(!fileDef || !fileDef.IsScalar()) EXCEPTION("Missing key \"Input file or directory\"");
  if(!plotDef || !plotDef.IsMap()) EXCEPTION("Missing key \"Plot definition\"");
  auto        file     = new BOSSOutputLoader(fileDef.as<string>().c_str());
  auto        hist     = file->Draw(plotDef);
  const char* namehist = Form("%s_hist", hist->GetName());
  auto        dataHist = new RooDataHist(namehist, namehist, fVar, Import(*hist));
  const char* namepdf  = Form("%s_pdf", hist->GetName());
  return new RooHistPdf(namepdf, namepdf, fVar, *dataHist, 0);
}

RooFFTConvPdf* RooFitBuilder::ConvolutePdf(RooAbsPdf* pdf1, RooAbsPdf* pdf2)
{
  if(!pdf1 || !pdf2) EXCEPTION("Nullptr pdfs");
  string name = Form("%s convoluted with %s", pdf1->GetName(), pdf2->GetName());
  return new RooFFTConvPdf(name.c_str(), name.c_str(), fVar, *pdf1, *pdf2);
}

RooAddPdf* RooFitBuilder::AddPdfs(RooArgList* pdfs, RooArgList* ratios)
{
  auto n = pdfs->getSize();
  if(n < 1) EXCEPTION("Cannot add fewer than one pdf");
  if(n - 1 != ratios->getSize() && n != ratios->getSize())
    EXCEPTION(Form("Ratios inconsistent with pdf list (%d | %d)", n, ratios->getSize()));
  string name = pdfs->at(0)->GetName();
  for(Int_t i = 0; i < n; ++i)
    name += Form("+%s", pdfs->at(i)->GetName());
  return new RooAddPdf(name.c_str(), name.c_str(), *pdfs, *ratios);
}

FitModelObj RooFitBuilder::BuildModel(const YAML::Node& input)
{
  const auto& sigTypeDef = input["type"];
  if(!sigTypeDef || !sigTypeDef.IsScalar())
    EXCEPTION("Missing key \"type\" in fit model definition");

  FitModelObj m;
  m.name = sigTypeDef.as<string>().c_str();
  try
  {
    m.factor = BuildVar(input["normalisation factor"]);
  }
  catch(const Exception& e)
  {}

  if(m.name.EqualTo("Histogram PDF"))
    m.pdf = BuildHistPdf(input);
  else if(m.name.EqualTo("Voigtian"))
    m.pdf = BuildVoigtian(input);
  else if(m.name.EqualTo("Breit-Wigner"))
    m.pdf = BuildBreitWigner(input);
  else if(m.name.EqualTo("Gaussian"))
    m.pdf = BuildGaussian(input);
  else if(m.name.EqualTo("P-wave"))
    m.pdf = BuildPwave(input);
  else if(m.name.EqualTo("Convoluted P-wave"))
    m.pdf = BuildConvolutedPwave(input);

  else if(m.name.EqualTo("Argus"))
    m.pdf = BuildArgus(input);
  else if(m.name.EqualTo("Argus BESIII"))
    m.pdf = BuildArgusBESIII(input);
  else if(m.name.EqualTo("Chebychev"))
    m.pdf = BuildChebychev(input);
  else if(m.name.EqualTo("Polynomial"))
    m.pdf = BuildPolynomial(input);
  else if(m.name.EqualTo("Truncated polynomial"))
    m.pdf = BuildTruncatedPolynomial(input);
  else
    EXCEPTION(Form("No fit model type defined for \"%s\"", m.name.Data()));

  return m;
}

SigPlusBck RooFitBuilder::BuildSigPlusBck(const YAML::Node& input)
{
  SigPlusBck  m;
  const auto& sigDef   = input["Signal"];
  const auto& bckDef   = input["Background"];
  const auto& ratioDef = input["Signal-to-background ratio"];
  if(!!sigDef && sigDef.IsMap()) m.sig = BuildModel(sigDef);
  if(!!bckDef && bckDef.IsMap()) m.bck = BuildModel(bckDef);

  m.sum = new RooArgList;
  if(m.sig.pdf) m.sum->add(*m.sig.pdf);
  if(m.bck.pdf) m.sum->add(*m.bck.pdf);

  m.coeff = new RooArgList;
  if(!ratioDef || !ratioDef.IsMap())
  {
    if(m.sig.factor) m.coeff->add(*m.sig.factor);
    if(m.bck.factor) m.coeff->add(*m.bck.factor);
  }
  else
  {
    m.ratio = BuildVar(ratioDef);
    m.coeff->add(*m.ratio);
  }
  m.pdf = AddPdfs(m.sum, m.coeff);
  return m;
}

string RooFitBuilder::ExtractName(const YAML::Node& input)
{
  if(!input) EXCEPTION("Empty YAML input");
  string name = Form("#it{p}_{%d}", fID);
  if(!!input["name"] && input["name"].IsScalar())
    name = input["name"].as<string>();
  else
    ++fID;
  return name;
}